import json
import os

import util.common
from util.llm import http_request

MODEL_NAME = os.getenv('MODEL_NAME', "deepseek-chat")

SYSTEM_PROMPT = """你是一名演练脚本专家，善于从设备、场景描述和应急预案中提取关键词用于搜索相关演练脚本文档。
我将给出设备、场景描述和应急预案信息，你需要理解这些信息，并从描述中提取关键词用于搜索。
提取的关键词可以是多个，关键词使用json array表示。
输出示例：
如果需要输出的内容为"网络安全事件"、"数据勒索"、"应急响应"，则返回的内容为：["网络安全事件","数据勒索","应急响应"]"""

PROMPT = """以下我将给出设备、场景描述和应急预案信息，请帮我提取关键词用于搜索演练脚本文档。结果用json array表示

设备：{}
场景描述：{}
应急预案：{}"""


async def query_rewrite(equipment: str, scene_desc: str, plan: dict):
    """
    使用大模型进行查询重写（原有方法）
    
    Args:
        equipment: 设备
        scene_desc: 场景描述
        plan: 应急预案
        
    Returns:
        关键词列表
    """
    print(f"🔍 query_rewrite开始处理")
    print(f"🔍 设备: {equipment}")
    print(f"🔍 场景描述: {scene_desc}")
    print(f"🔍 预案字段数: {len(plan)}")
    
    try:
        # 构建查询文本
        plan_text = json.dumps(plan, ensure_ascii=False, indent=2)
        query_text = f"设备：{equipment}，场景：{scene_desc}"
        
        result = await http_request(
            system_prompt=SYSTEM_PROMPT,
            history=[],
            prompt=PROMPT.format(equipment, scene_desc, plan_text),
            model=MODEL_NAME
        )
        print(f"✅ LLM查询重写成功，原始结果: {result}")
        
        result = util.common.remove_markdown_code_block(result)
        print(f"🔍 移除markdown代码块后: {result}")
        
        result = json.loads(result)
        print(f"🔍 JSON解析成功: {result}")
        
        # 添加原始查询关键词
        result.append(equipment)
        result.append(scene_desc)
        # 添加预案中的关键字段
        for key in plan.keys():
            if isinstance(key, str) and len(key) > 1:
                result.append(key)
        
        result = list(set(result))
        print(f"✅ 查询重写完成，最终结果: {result}")
        
    except Exception as e:
        print(f"❌ query_rewrite失败: {e}")
        import traceback
        traceback.print_exc()
        # 使用原始信息作为备选
        result = [equipment, scene_desc]
        for key in plan.keys():
            if isinstance(key, str) and len(key) > 1:
                result.append(key)
        result = list(set(result))
        print(f"🔍 使用原始信息作为备选: {result}")
    
    return result


async def query_rewrite_retrieval_only(equipment: str, scene_desc: str, plan: dict):
    """
    纯检索模式的查询重写，不使用大模型（新增方法）
    
    Args:
        equipment: 设备
        scene_desc: 场景描述
        plan: 应急预案
        
    Returns:
        关键词列表
    """
    print(f"🔍 query_rewrite_retrieval_only开始处理")
    print(f"🔍 设备: {equipment}")
    print(f"🔍 场景描述: {scene_desc}")
    print(f"🔍 预案字段数: {len(plan)}")
    
    try:
        # 直接基于输入信息提取关键词
        print(f"🔍 使用纯检索模式，直接关键词提取")
        
        # 基于输入信息直接提取关键词
        keywords = []
        
        # 添加原始信息
        keywords.append(equipment)
        keywords.append(scene_desc)
        
        # 从设备中提取关键词
        if "网络" in equipment:
            keywords.extend(["网络安全", "网络设备", "网络演练"])
        if "安全" in equipment:
            keywords.extend(["安全设备", "安全演练", "安全防护"])
        if "数据" in equipment:
            keywords.extend(["数据安全", "数据保护", "数据演练"])
        
        # 从场景描述中提取关键词
        if "演练" in scene_desc:
            keywords.extend(["应急演练", "演练脚本", "演练流程"])
        if "应急" in scene_desc:
            keywords.extend(["应急响应", "应急预案", "应急处置"])
        if "攻击" in scene_desc:
            keywords.extend(["网络攻击", "恶意攻击", "攻击演练"])
        if "泄露" in scene_desc:
            keywords.extend(["数据泄露", "信息泄露", "泄露演练"])
        
        # 从应急预案中提取关键词
        for key, value in plan.items():
            if isinstance(key, str) and len(key) > 1:
                keywords.append(key)
            if isinstance(value, str) and len(value) > 1:
                # 从值中提取关键词
                if "应急" in value:
                    keywords.extend(["应急响应", "应急预案"])
                if "安全" in value:
                    keywords.extend(["安全事件", "安全防护"])
                if "数据" in value:
                    keywords.extend(["数据安全", "数据保护"])
        
        # 去重
        keywords = list(set(keywords))
        
        print(f"✅ 纯检索模式关键词提取完成: {keywords}")
        
    except Exception as e:
        print(f"❌ query_rewrite_retrieval_only失败: {e}")
        import traceback
        traceback.print_exc()
        # 使用原始信息作为备选
        keywords = [equipment, scene_desc]
        for key in plan.keys():
            if isinstance(key, str) and len(key) > 1:
                keywords.append(key)
        keywords = list(set(keywords))
        print(f"🔍 使用原始信息作为备选: {keywords}")
    
    return keywords 